import os
import argparse
from tqdm import tqdm
import numpy as np
from pprint import pprint
import random
import matplotlib.pyplot as plt
import cv2
import sklearn.cluster

from data.dataset import build_dataset
from data.utils import read_image

def parse_args():
    parser = argparse.ArgumentParser(description='visualize 3d input distribution')
    parser.add_argument('--datasets', type=str, help='datasets name', required=True)
    parser.add_argument('-n', '--num', type=int, help='number of images', required=True)
    parser.add_argument('--cluster', type=int, help='number of clusters', default=256)
    args = parser.parse_args()
    return args

class TinyDataset:
    def __init__(self, dataset_names, max_size):
        self.dataset_dicts = build_dataset(*dataset_names)
        self.dataset_dicts = random.sample(self.dataset_dicts, max_size)

    def __len__(self):
        return len(self.dataset_dicts)

    def __getitem__(self, index):
        filename = self.dataset_dicts[index]['file_name']
        image = read_image(filename)
        return image

if __name__ == '__main__':
    args = parse_args()
    
    train_datasets = ["{}-train".format(dataset) for dataset in args.datasets.split(',')]
    dataset = TinyDataset(train_datasets, max_size=args.num)

    """ RGB space """
    ax_rgb = plt.subplot(222, projection='3d')
    ax_rg = plt.subplot(221)
    ax_gb = plt.subplot(223)
    ax_rb = plt.subplot(224)
    plt.subplots_adjust(wspace=0.3, hspace=0.3)

    pixels_list = [image.reshape(-1, image.shape[-1]) for image in tqdm(dataset) if image.shape[-1] == 3]
    pixels = np.concatenate(pixels_list)
    print("pixels shape:", pixels.shape)

    ax_rgb.scatter(pixels[:,0], pixels[:,1], pixels[:,2], s=0.5, color='b')
    ax_rg.scatter(pixels[:,0], pixels[:,1], s=0.5, color='b')
    ax_gb.scatter(pixels[:,1], pixels[:,2], s=0.5, color='b')
    ax_rb.scatter(pixels[:,0], pixels[:,2], s=0.5, color='b')

    print("clustering")

    kmeans = sklearn.cluster.KMeans(n_clusters=args.cluster, random_state=0, verbose=1).fit(pixels)
    cluster_centers = kmeans.cluster_centers_

    ax_rgb.scatter(cluster_centers[:,0], cluster_centers[:,1], cluster_centers[:,2], s=1, color='r')
    ax_rg.scatter(cluster_centers[:,0], cluster_centers[:,1], s=1, color='r')
    ax_gb.scatter(cluster_centers[:,1], cluster_centers[:,2], s=1, color='r')
    ax_rb.scatter(cluster_centers[:,0], cluster_centers[:,2], s=1, color='r')


    ax_rgb.set_xlabel('R')
    ax_rgb.set_ylabel('G')
    ax_rgb.set_zlabel('B')

    ax_rg.set_xlabel('R')
    ax_rg.set_ylabel('G')

    ax_gb.set_xlabel('G')
    ax_gb.set_ylabel('B')

    ax_rb.set_xlabel('R')
    ax_rb.set_ylabel('B')

    print("saving")

    output_path = os.path.join('output', 'cluster_3d_input', '&'.join(train_datasets), '3d_input_RGB.png')
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    plt.savefig(output_path, dpi=240)

    print("done")